package com.example.ai.rag;

import com.example.ai.embedding.SimilarSearchReq;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.DefaultChatClientBuilder;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.ai.rag.preretrieval.query.expansion.QueryExpander;
import org.springframework.ai.rag.preretrieval.query.transformation.RewriteQueryTransformer;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionTextParser;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.ResponseEntity;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;

@Slf4j
@RestController
@RequestMapping("/rag")
public class RAGController {
    @Autowired
    VectorStore vectorStore;
    @Autowired
    ChatModel chatModel;
    @PostMapping("/search/advisor/q_a")
    public ResponseEntity<?> searchWithQuestionAnswerAdvisor(@RequestBody SimilarSearchReq req) {
        SearchRequest searchRequest = SearchRequest.builder()
                .query(req.getContent())
                .topK(req.getTopK())
                .similarityThreshold(req.getThreshold())
                .filterExpression(req.getFilter())
                .build();
        QuestionAnswerAdvisor questionAnswerAdvisor = QuestionAnswerAdvisor.builder(vectorStore)
                .searchRequest(searchRequest)
                .build();
        ChatResponse response = ChatClient.builder(chatModel)
                .build().prompt()
                .advisors(questionAnswerAdvisor)
                .user(req.getContent())
                .call()
                .chatResponse();
        log.info("response:{}", response);
        return ResponseEntity.ok(response.getResult().getOutput().getText());
    }

    @PostMapping("/search/advisor/q_a2")
    public ResponseEntity<?> searchWithQuestionAnswerAdvisor2(@RequestBody SimilarSearchReq req) {

        SearchRequest searchRequest = SearchRequest.builder()
                .query(req.getContent())
                .topK(req.getTopK())
                .similarityThreshold(req.getThreshold())
                .build();
        QuestionAnswerAdvisor questionAnswerAdvisor = QuestionAnswerAdvisor.builder(vectorStore)
                .searchRequest(searchRequest)
                .build();
        ChatClient chatClient = ChatClient.builder(chatModel)
                .defaultAdvisors(questionAnswerAdvisor)
                .build();
        // Update filter expression at runtime
        String content = chatClient.prompt()
                .user(req.getContent())
                .advisors(a -> a.param(QuestionAnswerAdvisor.FILTER_EXPRESSION, req.getFilter()))
                .call()
                .content();
        log.info("content:{}", content);
        return ResponseEntity.ok(content);
    }
    @PostMapping("/search/advisor/raa")
    public ResponseEntity<?> retrievalAugmentationAdvisor(@RequestBody SimilarSearchReq req) {
        FilterExpressionTextParser parser = new FilterExpressionTextParser();
        Advisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .similarityThreshold(req.getThreshold())
                        .topK(req.getTopK())
                        .filterExpression(()->StringUtils.hasLength(req.getFilter())?parser.parse(req.getFilter()):null)
                        .vectorStore(vectorStore)
                        .build())
                .build();
        String content = ChatClient.builder(chatModel).build()
                .prompt()
                .advisors(retrievalAugmentationAdvisor)
                .user(req.getContent())
                .call()
                .content();

        log.info("content:{}", content);
        return ResponseEntity.ok(content);
    }

    @PostMapping("/search/advisor/raa2")
    public ResponseEntity<?> retrievalAugmentationAdvisor2(@RequestBody SimilarSearchReq req) {
        ChatClient.Builder  chatClientBuilder = ChatClient.builder(chatModel);
        RewriteQueryTransformer rewriteQueryTransformer = RewriteQueryTransformer.builder()
                .chatClientBuilder(chatClientBuilder.build().mutate())
                .build();

        QueryExpander queryExpander = MultiQueryExpander.builder()
                .chatClientBuilder(chatClientBuilder.build().mutate())
                .numberOfQueries(2)
                .build();
        FilterExpressionTextParser parser = new FilterExpressionTextParser();
        Advisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .similarityThreshold(req.getThreshold())
                        .topK(req.getTopK())
                        .filterExpression(()->StringUtils.hasLength(req.getFilter())?parser.parse(req.getFilter()):null)
                        .vectorStore(vectorStore)
                        .build())
                .queryTransformers(rewriteQueryTransformer)
                .queryExpander(queryExpander)
                .build();
        String content = ChatClient.builder(chatModel).build()
                .prompt()
                .advisors(retrievalAugmentationAdvisor)
                .user(req.getContent())
                .call()
                .content();

        log.info("content:{}", content);
        return ResponseEntity.ok(content);
    }

 }
